package com.zhupf.Backtracking;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
 * @author zhupf
 * @date 2024年01月11日 14:20
 * @Description
 *
 * 51. N 皇后
 */
public class SolveNQueens {

    private List<List<String>> ans;

    private Set<Integer> columns;
    private Set<Integer> diagonals1;
    private Set<Integer> diagonals2;


    public List<List<String>> solveNQueens(int n) {
        ans = new ArrayList<>();
        columns = new HashSet<>();
        diagonals1 = new HashSet<>();
        diagonals2 = new HashSet<>();
        List<String> res = new ArrayList<>();
        StringBuffer sb = new StringBuffer();
        for (int i = 0; i < n; i++) {
            sb.append(".");
        }
        dfs(0, n, res, sb.toString());
        return ans;
    }

    public void dfs(int idx, int n, List<String> res, String s) {
        if (idx == n) {
            ans.add(new ArrayList<>(res));
            return;
        }
        StringBuffer sb = new StringBuffer(s);
        for (int i = 0; i < n; i++) {
            if (ckeck(idx, i)) {
                sb.setCharAt(i, 'Q');
                columns.add(i);
                diagonals1.add(idx - i);
                diagonals2.add(idx + i);
                res.add(sb.toString());
                dfs(idx + 1, n, res, s);
                sb.setCharAt(i, '.');
                res.remove(res.size() - 1);
                columns.remove(i);
                diagonals1.remove(idx - i);
                diagonals2.remove(idx + i);
            }
        }
    }

    private boolean ckeck(int cow, int col) {
        if (columns.contains(col)) {
            return false;
        }
        if (diagonals1.contains(cow - col)) {
            return false;
        }
        if (diagonals2.contains(cow + col)) {
            return false;
        }
        return true;
    }
}
